//
// Created by syshen on 2021/4/17.
//

#include "pupilDetect.h"

bool pupilDetect::createInstance(std::string mnn_path, std::string landmark_mnn_path, std::string eyesegment_mnn_path, int size, int num_thread, float thresh, float nms) {
    detector = new Retinaface(mnn_path, size, thresh, nms, 4, false);
    facelandmark = new FaceLandMark(landmark_mnn_path, 128, 4); // FaceLandmark
    eyesegment = new EyeSegmentation(eyesegment_mnn_path, 128, 4);//
    if (detector && facelandmark && eyesegment)
        return true;
    else
        return false;
}

std::vector<std::vector<cv::Point>> pupilDetect::detect(cv::Mat &img, cv::Mat &left_mask, cv::Mat &right_mask) {
    std::vector<cv::Mat> rgbChannels(3);
    cv::split(img, rgbChannels);
    cv::Mat frame_gray = rgbChannels[2];
    std::vector<landmark> faces;
    detector->face_detect(img, faces);
    std::vector<std::vector<cv::Point>> points;

    for(auto face : faces){
        cv::Rect rect(int(face.x1), int(face.y1), int(face.x2-face.x1), int(face.y2-face.y1));

        std::vector<cv::Point> point = findEyes(frame_gray, face);

        std::vector<cv::Point> left_corner;
        std::vector<cv::Point> right_corner;
        facelandmark->landmark_detect(img, rect, left_corner, right_corner);

        cv::Rect left_rect = boundingRect(left_corner);
        cv::Rect right_rect = boundingRect(right_corner);

        cv::Point left_centroid;
        cv::Point right_centroid;
        left_mask = eyesegment->eye_segment(img, left_rect, left_centroid);
        right_mask = eyesegment->eye_segment(img, right_rect, right_centroid);

        for (size_t i = 0; i<point.size(); i++) {
            point[i].x += int(face.x1);
            point[i].y += int(face.y1);
        }

        points.push_back(point);
        points.push_back(left_corner);
        points.push_back(right_corner);

        std::vector<cv::Point> left_centroid_vec;
        std::vector<cv::Point> right_centroid_vec;
        left_centroid_vec.push_back(left_centroid);
        right_centroid_vec.push_back(right_centroid);
        points.push_back(left_centroid_vec);
        points.push_back(right_centroid_vec);
    }

    return points;
}

std::vector<std::vector<cv::Point>> pupilDetect::detect(cv::Mat &img) {
    std::vector<cv::Mat> rgbChannels(3);
    cv::split(img, rgbChannels);

    std::vector<landmark> faces;
    int flag = detector->face_detect(img, faces);
    if (flag < 0)
        return std::vector<std::vector<cv::Point>>{};
    std::vector<std::vector<cv::Point>> points;
    landmark face = faces[0];

    face.x1 = syshen_min(syshen_max(0, face.x1), img.cols-1);
    face.x2 = syshen_min(syshen_max(face.x1, face.x2), img.cols-1);
    face.y1 = syshen_min(syshen_max(0, face.y1), img.rows-1);
    face.y2 = syshen_min(syshen_max(face.y1, face.y2), img.rows-1);
    if (face.x2 <= face.x1 || face.y2 <= face.y1)
        return std::vector<std::vector<cv::Point>>{};

    cv::Rect rect(int(face.x1), int(face.y1), int(face.x2-face.x1), int(face.y2-face.y1));

    std::vector<cv::Point> left_corner;
    std::vector<cv::Point> right_corner;
    facelandmark->landmark_detect(img, rect, left_corner, right_corner);

    if (left_corner.size() <=0 || right_corner.size() <= 0)
        return std::vector<std::vector<cv::Point>>{};

    cv::Rect left_rect = boundingRect(left_corner);

    if (left_rect.width <= 0 || left_rect.height <= 0)
        return std::vector<std::vector<cv::Point>>{};
    cv::Point left_centroid;
    eyesegment->eye_segment(img, left_rect, left_centroid);
    points.emplace_back(std::vector<cv::Point>{left_corner[0], left_centroid, left_corner[left_corner.size()/2-1]});

//    cv::Point right_centroid;
//    cv::Rect right_rect = boundingRect(right_corner);
//    if (right_rect.width <= 0 || right_rect.height <= 0)
//        return std::vector<std::vector<cv::Point>>{};
//    eyesegment->eye_segment(img, right_rect, right_centroid);
//    points.emplace_back(std::vector<cv::Point>{right_corner[0], right_centroid, right_corner[right_corner.size()/2-1]});

    return points;
}

std::vector<std::vector<cv::Point>> pupilDetect::detect(cv::Mat img, cv::Mat &debugImage) {
    std::vector<cv::Mat> rgbChannels(3);
    cv::split(img, rgbChannels);
    cv::Mat frame_gray = rgbChannels[2];
    std::vector<landmark> faces;
    detector->face_detect(img, faces);
    std::vector<std::vector<cv::Point>> points;
    for( auto face : faces )
    {
        cv::Rect rect(int(face.x1), int(face.y1), int(face.x2-face.x1), int(face.y2-face.y1));
        rectangle(debugImage, rect, 1234);
    }

    for(auto face : faces){
        cv::Rect rect(int(face.x1), int(face.y1), int(face.x2-face.x1), int(face.y2-face.y1));
        std::vector<cv::Point> point = findEyes(frame_gray, rect);
        for (size_t i = 0; i<point.size(); i++) {
            point[i].x += int(face.x1);
            point[i].y += int(face.y1);
            cv::circle(debugImage, point[i], 1, (0,0,255), 1);
        }
        points.push_back(point);
    }
    return points;
}

std::vector<cv::Point> pupilDetect::findEyes(cv::Mat &frame_gray, landmark &landm, bool use_dynamic){
    cv::Rect face(int(landm.x1), int(landm.y1), int(landm.x2-landm.x1), int(landm.y2-landm.y1));
    cv::Mat faceROI = frame_gray(face);
    //cv::Mat debugFace = faceROI;

    int eye_region_width = int(face.width * (kEyePercentWidth/100.0f));
    int eye_region_height = int(face.width * (kEyePercentHeight/100.0f));
    int eye_region_top = int(face.height * (kEyePercentTop/100.0f));
    cv::Rect leftEyeRegion(int(face.width*(kEyePercentSide/100.0f)),
                           eye_region_top,eye_region_width,eye_region_height);
    cv::Rect rightEyeRegion(face.width - eye_region_width - int(face.width*(kEyePercentSide/100.0f)),
                            eye_region_top,eye_region_width,eye_region_height);

    //-- Find Eye Centers
    cv::Point leftPupil = findEyeCenter(faceROI,leftEyeRegion,"Left Eye");
    cv::Point rightPupil = findEyeCenter(faceROI,rightEyeRegion,"Right Eye");
    // get corner regions
    cv::Rect leftRightCornerRegion(leftEyeRegion);
    leftRightCornerRegion.width -= leftPupil.x;
    leftRightCornerRegion.x += leftPupil.x;
    leftRightCornerRegion.height /= 2;
    leftRightCornerRegion.y += leftRightCornerRegion.height / 2;
    cv::Rect leftLeftCornerRegion(leftEyeRegion);
    leftLeftCornerRegion.width = leftPupil.x;
    leftLeftCornerRegion.height /= 2;
    leftLeftCornerRegion.y += leftLeftCornerRegion.height / 2;
    /*cv::Rect rightLeftCornerRegion(rightEyeRegion);
    rightLeftCornerRegion.width = rightPupil.x;
    rightLeftCornerRegion.height /= 2;
    rightLeftCornerRegion.y += rightLeftCornerRegion.height / 2;
    cv::Rect rightRightCornerRegion(rightEyeRegion);
    rightRightCornerRegion.width -= rightPupil.x;
    rightRightCornerRegion.x += rightPupil.x;
    rightRightCornerRegion.height /= 2;
    rightRightCornerRegion.y += rightRightCornerRegion.height / 2;
    rectangle(debugFace,leftRightCornerRegion,200);
    rectangle(debugFace,leftLeftCornerRegion,200);
    rectangle(debugFace,rightLeftCornerRegion,200);
    rectangle(debugFace,rightRightCornerRegion,200);*/
    // change eye centers to face coordinates
    rightPupil.x += rightEyeRegion.x;
    rightPupil.y += rightEyeRegion.y;
    leftPupil.x += leftEyeRegion.x;
    leftPupil.y += leftEyeRegion.y;

    location left, right;
    left._x = landm.point[0]._x - face.x;
    left._y = landm.point[0]._y - face.y;
    right._x = landm.point[1]._x - face.x;
    right._y = landm.point[1]._y - face.y;
    float abs_l_x = left._x > leftPupil.x ? left._x - leftPupil.x : leftPupil.x - left._x;
    float abs_l_y = left._y > leftPupil.y ? left._y - leftPupil.y : leftPupil.y - left._y;
    float abs_r_x = right._x > rightEyeRegion.x ? right._x - rightEyeRegion.x : rightEyeRegion.x - right._x;
    float abs_r_y = right._y > rightEyeRegion.y ? right._y - rightEyeRegion.y : rightEyeRegion.y - right._y;

    if(abs_l_x > 10 || abs_l_y > 10) {
        leftPupil = cv::Point(int(left._x), int(left._y));
    }
    if(abs_r_x > 10 || abs_r_y > 10) {
        rightPupil = cv::Point(int(right._x), int(right._y));
    }

    return std::vector<cv::Point>{leftPupil, rightPupil};
}

std::vector<cv::Point> pupilDetect::findEyes(cv::Mat &frame_gray, cv::Rect face){
    cv::Mat faceROI = frame_gray(face);
//        cv::Mat debugFace = faceROI;

    //-- Find eye regions and draw them
    int eye_region_width = int(face.width * (kEyePercentWidth/100.0f));
    int eye_region_height = int(face.width * (kEyePercentHeight/100.0f));
    int eye_region_top = int(face.height * (kEyePercentTop/100.0f));
    cv::Rect leftEyeRegion(int(face.width*(kEyePercentSide/100.0f)),
                           eye_region_top,eye_region_width,eye_region_height);
    cv::Rect rightEyeRegion(face.width - eye_region_width - int(face.width*(kEyePercentSide/100.0f)),
                            eye_region_top,eye_region_width,eye_region_height);

    //-- Find Eye Centers
    cv::Point leftPupil = findEyeCenter(faceROI,leftEyeRegion,"Left Eye");
    cv::Point rightPupil = findEyeCenter(faceROI,rightEyeRegion,"Right Eye");
    // get corner regions
    cv::Rect leftRightCornerRegion(leftEyeRegion);
    leftRightCornerRegion.width -= leftPupil.x;
    leftRightCornerRegion.x += leftPupil.x;
    leftRightCornerRegion.height /= 2;
    leftRightCornerRegion.y += leftRightCornerRegion.height / 2;
    cv::Rect leftLeftCornerRegion(leftEyeRegion);
    leftLeftCornerRegion.width = leftPupil.x;
    leftLeftCornerRegion.height /= 2;
    leftLeftCornerRegion.y += leftLeftCornerRegion.height / 2;
    // change eye centers to face coordinates
    rightPupil.x += rightEyeRegion.x;
    rightPupil.y += rightEyeRegion.y;
    leftPupil.x += leftEyeRegion.x;
    leftPupil.y += leftEyeRegion.y;

    return std::vector<cv::Point>{leftPupil, rightPupil};
}

std::vector<cv::Point> pupilDetect::findEyeCorners(cv::Mat &frame_gray, landmark &landm, cv::Point leftPupil,
                                      cv::Point rightPupil){
    cv::Rect face(int(landm.x1), int(landm.y1), int(landm.x2-landm.x1), int(landm.y2-landm.y1));
    cv::Mat faceROI = frame_gray(face);
    //-- Find eye regions and draw them
    int eye_region_width = int(face.width * (kEyePercentWidth/100.0f));
    int eye_region_height = int(face.width * (kEyePercentHeight/100.0f));
    int eye_region_top = int(face.height * (kEyePercentTop/100.0f));
    cv::Rect leftEyeRegion(int(face.width*(kEyePercentSide/100.0f)),
                           eye_region_top,eye_region_width,eye_region_height);
    cv::Rect rightEyeRegion(face.width - eye_region_width - int(face.width*(kEyePercentSide/100.0f)),
                            eye_region_top,eye_region_width,eye_region_height);
    cv::Rect leftRightCornerRegion(leftEyeRegion);
    leftRightCornerRegion.width -= leftPupil.x;
    leftRightCornerRegion.x += leftPupil.x;
    leftRightCornerRegion.height /= 2;
    leftRightCornerRegion.y += leftRightCornerRegion.height / 2;
    cv::Rect leftLeftCornerRegion(leftEyeRegion);
    leftLeftCornerRegion.width = leftPupil.x;
    leftLeftCornerRegion.height /= 2;
    leftLeftCornerRegion.y += leftLeftCornerRegion.height / 2;
    cv::Rect rightLeftCornerRegion(rightEyeRegion);
    rightLeftCornerRegion.width = rightPupil.x;
    rightLeftCornerRegion.height /= 2;
    rightLeftCornerRegion.y += rightLeftCornerRegion.height / 2;
    cv::Rect rightRightCornerRegion(rightEyeRegion);
    rightRightCornerRegion.width -= rightPupil.x;
    rightRightCornerRegion.x += rightPupil.x;
    rightRightCornerRegion.height /= 2;
    rightRightCornerRegion.y += rightRightCornerRegion.height / 2;
    cv::Point2f leftRightCorner = findEyeCorner(faceROI(leftRightCornerRegion), true, false);
    leftRightCorner.x += leftRightCornerRegion.x;
    leftRightCorner.y += leftRightCornerRegion.y;
    cv::Point2f leftLeftCorner = findEyeCorner(faceROI(leftLeftCornerRegion), true, true);
    leftLeftCorner.x += leftLeftCornerRegion.x;
    leftLeftCorner.y += leftLeftCornerRegion.y;
    cv::Point2f rightLeftCorner = findEyeCorner(faceROI(rightLeftCornerRegion), false, true);
    rightLeftCorner.x += rightLeftCornerRegion.x;
    rightLeftCorner.y += rightLeftCornerRegion.y;
    cv::Point2f rightRightCorner = findEyeCorner(faceROI(rightRightCornerRegion), false, false);
    rightRightCorner.x += rightRightCornerRegion.x;
    rightRightCorner.y += rightRightCornerRegion.y;

    return std::vector<cv::Point>{leftLeftCorner, rightRightCorner};
}

std::vector<cv::Point> pupilDetect::findEyeCorners(cv::Mat &frame_gray, cv::Rect &face, cv::Point &leftPupil,
                                      cv::Point &rightPupil){
    cv::Mat faceROI = frame_gray(face);
    //-- Find eye regions and draw them
    int eye_region_width = int(face.width * (kEyePercentWidth/100.0f));
    int eye_region_height = int(face.width * (kEyePercentHeight/100.0f));
    int eye_region_top = int(face.height * (kEyePercentTop/100.0f));
    cv::Rect leftEyeRegion(int(face.width*(kEyePercentSide/100.0f)),
                           eye_region_top,eye_region_width,eye_region_height);
    cv::Rect rightEyeRegion(face.width - eye_region_width - int(face.width*(kEyePercentSide/100.0f)),
                            eye_region_top,eye_region_width,eye_region_height);
    cv::Rect leftRightCornerRegion(leftEyeRegion);
    leftRightCornerRegion.width -= leftPupil.x;
    leftRightCornerRegion.x += leftPupil.x;
    leftRightCornerRegion.height /= 2;
    leftRightCornerRegion.y += leftRightCornerRegion.height / 2;
    cv::Rect leftLeftCornerRegion(leftEyeRegion);
    leftLeftCornerRegion.width = leftPupil.x;
    leftLeftCornerRegion.height /= 2;
    leftLeftCornerRegion.y += leftLeftCornerRegion.height / 2;
    cv::Rect rightLeftCornerRegion(rightEyeRegion);
    rightLeftCornerRegion.width = rightPupil.x;
    rightLeftCornerRegion.height /= 2;
    rightLeftCornerRegion.y += rightLeftCornerRegion.height / 2;
    cv::Rect rightRightCornerRegion(rightEyeRegion);
    rightRightCornerRegion.width -= rightPupil.x;
    rightRightCornerRegion.x += rightPupil.x;
    rightRightCornerRegion.height /= 2;
    rightRightCornerRegion.y += rightRightCornerRegion.height / 2;
    cv::Point2f leftRightCorner = findEyeCorner(faceROI(leftRightCornerRegion), true, false);
    leftRightCorner.x += leftRightCornerRegion.x;
    leftRightCorner.y += leftRightCornerRegion.y;
    cv::Point2f leftLeftCorner = findEyeCorner(faceROI(leftLeftCornerRegion), true, true);
    leftLeftCorner.x += leftLeftCornerRegion.x;
    leftLeftCorner.y += leftLeftCornerRegion.y;
    cv::Point2f rightLeftCorner = findEyeCorner(faceROI(rightLeftCornerRegion), false, true);
    rightLeftCorner.x += rightLeftCornerRegion.x;
    rightLeftCorner.y += rightLeftCornerRegion.y;
    cv::Point2f rightRightCorner = findEyeCorner(faceROI(rightRightCornerRegion), false, false);
    rightRightCorner.x += rightRightCornerRegion.x;
    rightRightCorner.y += rightRightCornerRegion.y;
    return std::vector<cv::Point>{leftLeftCorner, rightRightCorner};
}
